Skip to content

returning model raw outputs as well - #8

Open
arthrod wants to merge 1 commit into
cicero-im:mainfrom
ajay-sreeram:main
Open

returning model raw outputs as well#8
arthrod wants to merge 1 commit into
cicero-im:mainfrom
ajay-sreeram:main

Conversation

@arthrod

@arthrod arthrod commented Feb 17, 2026

Copy link
Copy Markdown
Collaborator

Summary by Sourcery

Return both extracted entities and raw model outputs from prediction helpers.

New Features:

  • Expose raw model outputs alongside extracted entities in batch entity prediction API.

Enhancements:

  • Adjust single-text prediction helper to work with the updated batch prediction return structure without changing its public return type.

@sourcery-ai

sourcery-ai Bot commented Feb 17, 2026

Copy link
Copy Markdown
Reviewer's guide (collapsed on small PRs)

Reviewer's Guide

This PR updates the entity prediction APIs to also expose the raw model outputs while preserving the previous return shape for single-text predictions, and refactors the internal call to keep the model output object intact before indexing into its first element.

Sequence diagram for updated entity prediction and raw outputs

sequenceDiagram
    actor Client
    participant GlinerModel
    participant UnderlyingModel

    Client->>GlinerModel: predict_entities(text, labels, flat_ner, threshold, multi_label)
    GlinerModel->>GlinerModel: batch_predict_entities([text], labels, flat_ner, threshold, multi_label)
    GlinerModel->>UnderlyingModel: model(model_input)
    UnderlyingModel-->>GlinerModel: model_output_raw
    GlinerModel->>GlinerModel: model_output = model_output_raw[0]
    GlinerModel->>GlinerModel: postprocess to all_entities
    GlinerModel-->>GlinerModel: return all_entities, model_output_raw
    GlinerModel->>GlinerModel: all_entities[0]
    GlinerModel-->>Client: entities_for_single_text
Loading

Class diagram for updated prediction methods exposing raw outputs

classDiagram
    class GlinerModel {
        +predict_entities(text, labels, flat_ner, threshold, multi_label) entities
        +batch_predict_entities(texts, labels, flat_ner, threshold, multi_label) (all_entities, model_output_raw)
        +prepare_model_inputs(texts, labels) (model_input, raw_batch)
    }

    class UnderlyingModel {
        +__call__(model_input) model_output_raw
    }

    GlinerModel --> UnderlyingModel : uses for inference
Loading

File-Level Changes

Change Details Files
Expose raw model outputs from batch entity prediction and adapt single-text prediction accordingly.
  • Change predict_entities to unpack the tuple returned by batch_predict_entities and continue returning only the first entry of the entities list to preserve its original public API
  • Update batch_predict_entities to call the underlying model once, keep the full returned object as model_output_raw, and derive the tensor logits from its first element
  • Adjust batch_predict_entities return value from just the list of all_entities to a tuple of (all_entities, model_output_raw) so callers can access both processed entities and raw model outputs
gliner/model.py

Tips and commands

Interacting with Sourcery

  • Trigger a new review: Comment @sourcery-ai review on the pull request.
  • Continue discussions: Reply directly to Sourcery's review comments.
  • Generate a GitHub issue from a review comment: Ask Sourcery to create an
    issue from a review comment by replying to it. You can also reply to a
    review comment with @sourcery-ai issue to create an issue from it.
  • Generate a pull request title: Write @sourcery-ai anywhere in the pull
    request title to generate a title at any time. You can also comment
    @sourcery-ai title on the pull request to (re-)generate the title at any time.
  • Generate a pull request summary: Write @sourcery-ai summary anywhere in
    the pull request body to generate a PR summary at any time exactly where you
    want it. You can also comment @sourcery-ai summary on the pull request to
    (re-)generate the summary at any time.
  • Generate reviewer's guide: Comment @sourcery-ai guide on the pull
    request to (re-)generate the reviewer's guide at any time.
  • Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
    pull request to resolve all Sourcery comments. Useful if you've already
    addressed all the comments and don't want to see them anymore.
  • Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
    request to dismiss all existing Sourcery reviews. Especially useful if you
    want to start fresh with a new review - don't forget to comment
    @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

  • Enable or disable review features such as the Sourcery-generated pull request
    summary, the reviewer's guide, and others.
  • Change the review language.
  • Add, remove or edit custom review instructions.
  • Adjust other review settings.

Getting Help

@gemini-code-assist

Copy link
Copy Markdown

Summary of Changes

Hello @arthrod, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the GLiNER model's prediction capabilities by exposing the raw outputs from the underlying model. This allows users to access more detailed information beyond just the extracted entities, which can be valuable for debugging, advanced analysis, or custom post-processing. The primary predict_entities method retains its original output format for backward compatibility, while the batch prediction method now provides the additional raw data.

Highlights

  • Raw Model Output Exposure: The batch_predict_entities method now returns the raw outputs from the underlying model in addition to the processed entities. This provides more detailed information for advanced use cases.
  • API Adaptation: The predict_entities method has been updated to accommodate the new return signature of batch_predict_entities, ensuring that single entity predictions continue to return only the processed entities while internally handling the additional raw output.
Changelog
  • gliner/model.py
    • Modified batch_predict_entities to return both processed entities and raw model outputs.
    • Updated predict_entities to handle the new return signature of batch_predict_entities.
    • Captured the full raw output from the internal model call within batch_predict_entities.
    • Ensured the file ends with a newline character.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@coderabbitai

coderabbitai Bot commented Feb 17, 2026

Copy link
Copy Markdown
📝 Walkthrough

Summary by CodeRabbit

  • New Features

    • Entity prediction methods now provide access to raw model outputs alongside decoded entity results, enabling advanced use cases requiring low-level model data.
  • Breaking Changes

    • Batch entity prediction return format updated to include both entity predictions and raw model outputs as a tuple.

Walkthrough

The predict_entities and batch_predict_entities methods in the GLiNER model class have been refactored to return tuples containing both entity predictions and raw model output, instead of only returning entities. This provides access to underlying ONNX/torch model outputs.

Changes

Cohort / File(s) Summary
GLiNER Model Methods
gliner/model.py
Modified batch_predict_entities to return a tuple (all_entities, model_output_raw) instead of just all_entities. Updated predict_entities to unpack this tuple and return the first element of the decoded entities. Raw model output is now captured and exposed alongside entity predictions.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Poem

🐰 Two methods dance in tuple's grace,
Raw output finds its rightful place,
No longer hidden from our sight,
The model's whispers, pure delight!
With entities and tensors paired,
A richer truth is now declared.

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Merge Conflict Detection ⚠️ Warning ❌ Merge conflicts detected (33 files):

⚔️ .gitignore (content)
⚔️ README.md (content)
⚔️ RELEASE.md (content)
⚔️ examples/convert_to_onnx.ipynb (content)
⚔️ examples/finetune.ipynb (content)
⚔️ examples/load_local_model.ipynb (content)
⚔️ examples/quickstart.ipynb (content)
⚔️ gliner/__init__.py (content)
⚔️ gliner/config.py (content)
⚔️ gliner/data_processing/__init__.py (content)
⚔️ gliner/data_processing/collator.py (content)
⚔️ gliner/data_processing/processor.py (content)
⚔️ gliner/data_processing/tokenizer.py (content)
⚔️ gliner/data_processing/utils.py (content)
⚔️ gliner/decoding/__init__.py (content)
⚔️ gliner/decoding/decoder.py (content)
⚔️ gliner/decoding/utils.py (content)
⚔️ gliner/evaluation/__init__.py (content)
⚔️ gliner/evaluation/evaluator.py (content)
⚔️ gliner/model.py (content)
⚔️ gliner/modeling/base.py (content)
⚔️ gliner/modeling/encoder.py (content)
⚔️ gliner/modeling/layers.py (content)
⚔️ gliner/modeling/loss_functions.py (content)
⚔️ gliner/modeling/scorers.py (content)
⚔️ gliner/modeling/span_rep.py (content)
⚔️ gliner/onnx/model.py (content)
⚔️ gliner/training/__init__.py (content)
⚔️ gliner/training/trainer.py (content)
⚔️ gliner/utils.py (content)
⚔️ pyproject.toml (content)
⚔️ requirements.txt (content)
⚔️ train.py (content)

These conflicts must be resolved before merging into main.
Resolve conflicts locally and push changes to this branch.
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title 'returning model raw outputs as well' directly describes the main change of exposing raw model outputs alongside extracted entities, matching the primary objective of the PR.
Description check ✅ Passed The description clearly relates to the changeset, explaining both the new feature (raw model outputs) and enhancements (updated batch prediction API) reflected in the code changes.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
⚔️ Resolve merge conflicts (beta)
  • Auto-commit resolved conflicts to branch main
  • Post resolved changes as copyable diffs in a comment

Comment @coderabbitai help to get the list of available commands and usage tips.

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey - I've left some high level feedback:

  • Changing batch_predict_entities to return (all_entities, model_output_raw) instead of just all_entities is a breaking API change; consider either adding a new method for the extended return value or keeping the old return shape and exposing the raw outputs via an optional flag or attribute to avoid breaking existing callers.
  • In predict_entities, model_output is assigned but never used and the method still only returns all_entities[0]; either return the raw outputs there as well or remove the unused variable to keep the intent clear.
  • Update the docstrings (or function comments) for predict_entities and batch_predict_entities to describe the new return values so callers understand that the batch method now returns both entities and raw model outputs.
Prompt for AI Agents
Please address the comments from this code review:

## Overall Comments
- Changing `batch_predict_entities` to return `(all_entities, model_output_raw)` instead of just `all_entities` is a breaking API change; consider either adding a new method for the extended return value or keeping the old return shape and exposing the raw outputs via an optional flag or attribute to avoid breaking existing callers.
- In `predict_entities`, `model_output` is assigned but never used and the method still only returns `all_entities[0]`; either return the raw outputs there as well or remove the unused variable to keep the intent clear.
- Update the docstrings (or function comments) for `predict_entities` and `batch_predict_entities` to describe the new return values so callers understand that the batch method now returns both entities and raw model outputs.

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b6df740f6f

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread gliner/model.py
all_entities.append(entities)

return all_entities
return all_entities, model_output_raw

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve batch_predict_entities return contract

Changing batch_predict_entities to always return (all_entities, model_output_raw) breaks existing callers that were written against the previous contract (entities = model.batch_predict_entities(...) returning only entity lists). In downstream inference pipelines this can immediately raise type errors or silently corrupt post-processing when code iterates/serializes the returned value, so this should be made backward-compatible (for example via an opt-in flag for raw outputs) rather than a hard return-type change.

Useful? React with 👍 / 👎.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
gliner/model.py (1)

138-174: ⚠️ Potential issue | 🟠 Major

Breaking API change: update docstring and consider backward compatibility.

Changing the return type from List to Tuple[List, Any] is a breaking change. Existing callers using:

entities = model.batch_predict_entities(texts, labels)

will now receive a tuple and likely break.

Two issues to address:

  1. Update the docstring to document the new return type:
         """
         Predict entities for a batch of texts.
         texts:  List of texts | List[str]
         labels: List of labels | List[str]
         ...
+        
+        Returns:
+            Tuple[List[List[Dict]], Any]: A tuple containing (all_entities, model_output_raw)
         """
  1. Consider backward compatibility via an optional parameter:
def batch_predict_entities(self, texts, labels, flat_ner=True, threshold=0.5, 
                           multi_label=False, return_model_output=False):
    ...
    if return_model_output:
        return all_entities, model_output_raw
    return all_entities
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@gliner/model.py` around lines 138 - 174, The batch_predict_entities function
now returns a tuple (all_entities, model_output_raw) which is a breaking API
change; update the function docstring to document the new return value and add a
backward-compatible optional flag (e.g., return_model_output=False) to preserve
the original behavior. Modify the signature of batch_predict_entities to accept
return_model_output=False, and at the end return (all_entities,
model_output_raw) only when return_model_output is True, otherwise return just
all_entities; reference the existing symbols model_output_raw, all_entities, and
the function name batch_predict_entities when making these changes.
🧹 Nitpick comments (1)
gliner/model.py (1)

132-135: Unused variable and API inconsistency.

The model_output variable is unpacked but never used. Per Python convention, prefix it with an underscore to indicate it's intentionally discarded:

-        all_entities, model_output = self.batch_predict_entities(
+        all_entities, _model_output = self.batch_predict_entities(
             [text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label
         )

Additionally, should predict_entities also return raw outputs for API consistency with batch_predict_entities? If backward compatibility is the goal, consider adding an optional return_raw_output=False parameter instead.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@gliner/model.py` around lines 132 - 135, The local variable model_output in
predict_entities is unused — rename it to _model_output when unpacking the
result of batch_predict_entities to signal it's intentionally discarded (i.e.,
replace model_output with _model_output in predict_entities). If you want API
consistency with batch_predict_entities, add an optional parameter
return_raw_output: bool = False to predict_entities' signature and, when True,
return a tuple (all_entities[0], _model_output) otherwise keep the current
behavior of returning all_entities[0]; ensure the call to batch_predict_entities
preserves the existing parameters (labels, flat_ner, threshold, multi_label) and
maintain backward compatibility by defaulting return_raw_output to False.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Outside diff comments:
In `@gliner/model.py`:
- Around line 138-174: The batch_predict_entities function now returns a tuple
(all_entities, model_output_raw) which is a breaking API change; update the
function docstring to document the new return value and add a
backward-compatible optional flag (e.g., return_model_output=False) to preserve
the original behavior. Modify the signature of batch_predict_entities to accept
return_model_output=False, and at the end return (all_entities,
model_output_raw) only when return_model_output is True, otherwise return just
all_entities; reference the existing symbols model_output_raw, all_entities, and
the function name batch_predict_entities when making these changes.

---

Nitpick comments:
In `@gliner/model.py`:
- Around line 132-135: The local variable model_output in predict_entities is
unused — rename it to _model_output when unpacking the result of
batch_predict_entities to signal it's intentionally discarded (i.e., replace
model_output with _model_output in predict_entities). If you want API
consistency with batch_predict_entities, add an optional parameter
return_raw_output: bool = False to predict_entities' signature and, when True,
return a tuple (all_entities[0], _model_output) otherwise keep the current
behavior of returning all_entities[0]; ensure the call to batch_predict_entities
preserves the existing parameters (labels, flat_ner, threshold, multi_label) and
maintain backward compatibility by defaulting return_raw_output to False.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request modifies the prediction helpers to return raw model outputs, updating batch_predict_entities and predict_entities accordingly. It also resolves a critical issue where the original model response was an empty, invalid JSON, by ensuring a minimal valid JSON structure is always returned. The changes are clear and achieve the stated goal, with a minor suggestion for improving code clarity by handling an unused variable.

Comment thread gliner/model.py
Comment on lines +132 to +134
all_entities, model_output = self.batch_predict_entities(
[text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label
)[0]
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The model_output variable is assigned but never used. To signal that it's intentionally being ignored, you can use an underscore _ as the variable name. This improves code clarity and adheres to common Python conventions for unused variables.

Suggested change
all_entities, model_output = self.batch_predict_entities(
[text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label
)[0]
)
all_entities, _ = self.batch_predict_entities(
[text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label
)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants